import cv2
import numpy as np
import random
import os


def augment_image(image):
    """对单张图像进行增强操作"""
    # 获取图片的高度和宽度
    h, w = image.shape[:2]

    # 随机缩放（放大或缩小）
    scale = random.uniform(0.8, 1.2)  # 0.8 到 1.2 的随机缩放因子
    new_size = (int(w * scale), int(h * scale))
    image = cv2.resize(image, new_size)

    # 随机旋转
    angle = random.uniform(0, 30)  # 0 到 30 度的随机角度
    center = (new_size[0] // 2, new_size[1] // 2)
    rotation_matrix = cv2.getRotationMatrix2D(center, angle, 1.0)
    image = cv2.warpAffine(image, rotation_matrix, new_size)

    # 50% 概率进行水平翻转
    if random.random() < 0.5:
        image = cv2.flip(image, 1)  # 1 表示水平翻转

    # 50% 概率随机应用高斯滤波
    if random.random() < 0.5:
        kernel_size = random.choice([3, 5, 7])  # 随机选择高斯核大小
        image = cv2.GaussianBlur(image, (kernel_size, kernel_size), 0)

    return image


def process_dataset(dataset_path, output_path):
    """处理数据集中的所有图像，保存增强后的图像"""
    os.makedirs(output_path, exist_ok=True)

    for filename in os.listdir(dataset_path):
        image_path = os.path.join(dataset_path, filename)
        image = cv2.imread(image_path)

        # 30% 概率决定是否增强图像
        if random.random() < 0.3:
            enhanced_image = augment_image(image)

            # 保存增强后的图片
            base_name, ext = os.path.splitext(filename)
            enhanced_image_path = os.path.join(output_path, f"{base_name}_enhanced{ext}")
            cv2.imwrite(enhanced_image_path, enhanced_image)
            print(f"Enhanced image saved as: {enhanced_image_path}")


if __name__ == "__main__":
    # dataset_path = 'path/to/your/dataset'  # 替换为数据集路径
    # output_path = 'path/to/save/enhanced_images'  # 替换为输出路径

    dataset_path = 'VOC2028_short/images1'  # 替换为原始数据集路径
    output_path = 'VOC2028_short/images3'  # 替换为输出路径

    process_dataset(dataset_path, output_path)